from gpt2_download import download_and_load_gpt2  # 从gpt_download导入download_and_load_gpt2函数
settings, params = download_and_load_gpt2(model_size="124M", models_dir="gpt2")  # 下载并加载GPT-2设置和参数
from gpt2_tools import *

GPT_CONFIG_124M = {
    "vocab_size": 50257,
  # 词汇量大小
    "context_length": 256,  # 上下文长度
    "emb_dim": 768,  # 嵌入维度
    "n_heads": 12,  # 注意力头数
    "n_layers": 12,  # 层数
    "drop_rate": 0.1,  # 丢弃率
    "qkv_bias": False  # 查询-键-值偏差
}

model_configs = {  # 定义模型配置字典
    "gpt2-small (124M)": {"emb_dim": 768, "n_layers": 12, "n_heads": 12},
    "gpt2-medium (355M)": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
    "gpt2-large (774M)": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
    "gpt2-xl (1558M)": {"emb_dim": 1600, "n_layers": 48, "n_heads": 25},
}

file_path = "the-verdict.txt"  # 文件路径
with open(file_path, "r", encoding="utf-8") as file:  # 以读模式打开文件
    text_data = file.read()  # 读取文件内容

train_ratio = 0.90  # 训练集比例
split_idx = int(train_ratio * len(text_data))  # 计算分割索引
train_data = text_data[:split_idx]  # 获取训练数据
val_data = text_data[split_idx:]  # 获取验证数据

train_loader = create_dataloader_v1(
    train_data,  # 训练数据
    batch_size=2,  # 批大小
    max_length=GPT_CONFIG_124M["context_length"],  # 最大长度
    stride=GPT_CONFIG_124M["context_length"],  # 步幅
    drop_last=True,  # 丢弃最后一个不完整批次
    shuffle=True,  # 是否打乱数据
    num_workers=0  # 工作线程数
)

val_loader = create_dataloader_v1(
    val_data,  # 验证数据
    batch_size=2,  # 批大小
    max_length=GPT_CONFIG_124M["context_length"],  # 最大长度
    stride=GPT_CONFIG_124M["context_length"],  # 步幅
    drop_last=False,  # 不丢弃最后一个不完整批次
    shuffle=False,  # 是否打乱数据
    num_workers=0  # 工作线程数
)

tokenizer = tiktoken.get_encoding("gpt2") # 获取GPT-2的分词器编码
torch.manual_seed(123)              # 设置随机种子
model_name = "gpt2-small (124M)"  # 选择模型名称
NEW_CONFIG = GPT_CONFIG_124M.copy()  # 复制原始配置
NEW_CONFIG.update(model_configs[model_name])  # 更新配置为选定模型的配置
NEW_CONFIG.update({"context_length": 1024})  # 更新上下文长度为1024
NEW_CONFIG.update({"qkv_bias": True})  # 启用偏置向量
gpt = GPTModel(NEW_CONFIG)   # 使用配置初始化模型
load_weights_into_gpt(gpt, params)  # 加载权重到GPT模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 设置设备为GPU或CPU
gpt.to(device)  # 将模型移动到设备上
optimizer = torch.optim.AdamW(gpt.parameters(), lr=0.0004, weight_decay=0.1)  # 使用AdamW优化器
num_epochs = 10  # 训练周期数
train_losses, val_losses, tokens_seen = train_model_simple(  # 调用train_model_simple函数
    gpt, train_loader, val_loader, optimizer, device,
    num_epochs=num_epochs, eval_freq=5, eval_iter=1,
    start_context="Every effort moves you", tokenizer=tokenizer
)

epochs_tensor = torch.linspace(0, num_epochs, len(train_losses))  # 创建一个从0到num_epochs的线性张量，长度为训练损失的长度
plot_losses(epochs_tensor, tokens_seen, train_losses,
            val_losses)  # 调用plot_losses函数，传入epochs_tensor, tokens_seen, train_losses, val_losses

torch.save({  # 保存模型和优化器的状态字典
    "model_state_dict": gpt.state_dict(),  # 模型状态字典
    "optimizer_state_dict": optimizer.state_dict(),  # 优化器状态字典
}, "model_and_optimizer.pth")  # 保存到model_and_optimizer.pth文件

# 测试效果
token_ids = generate(
    model=gpt,
    idx=text_to_token_ids("Every effort moves you", tokenizer).to(device),  # 将文本转换为词元ID并移动到设备
    max_new_tokens=15,  # 最大新词元数为15
    context_size=GPT_CONFIG_124M["context_length"],  # 上下文大小
    top_k=25,  # top-k值为25
    temperature=1.4  # 温度值为1.4
)
print("Output text:\n", token_ids_to_text(token_ids, tokenizer))  # 打印生成的文本

# 后续可以加载之前训练过程中保存的模型和优化器状态继续训练
# checkpoint = torch.load("model_and_optimizer.pth")  # 加载保存的检查点
# model = GPTModel(GPT_CONFIG_124M)  # 初始化新模型实例
# model.load_state_dict(checkpoint["model_state_dict"])  # 加载模型状态字典
# optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.1)  # 初始化优化器
# optimizer.load_state_dict(checkpoint["optimizer_state_dict"])  # 加载优化器状态字典
# model.train()  # 设置模型为训练模式







